#! /usr/bin/env python3
# -*- coding: utf-8-*

import rospy
from sensor_msgs.msg import Image
from block_detection.srv import *
from cv_bridge import CvBridge
import cv2 as cv
import numpy as np


class Detection:
    def __init__(self):
        # 角点值
        self.corner_point = np.float32([[18, 50], [573, 14], [610, 433], [40, 471]])
        # 计算变换矩阵
        pst_obs = np.float32([[0, 0], [400, 0], [400, 300], [0, 300]])
        self.Mat = cv.getPerspectiveTransform(self.corner_point, pst_obs)   

        print(cv.__version__)
        

    def ros_running(self):
        rospy.init_node("block_detection")
        camera_name = rospy.get_param("~camera_name", '')
        self.knn_model = rospy.get_param("~knn_model",'')
        knn = cv.ml.KNearest_create()
        self.model = knn.load(self.knn_model)
        rospy.Subscriber(camera_name,Image,self.image_callback)
        rospy.Service('block_detection',block_result,self.detection_callback)
        rospy.spin()


    def image_callback(self, frame):
        bridge = CvBridge()
        # 消息转换为rgb uint8
        self.frame = bridge.imgmsg_to_cv2(frame, 'bgr8')
        # rec = self.predict(self.frame)
        # print(rec)


    def detection_callback(self, req):
        if req:
            return block_resultResponse(self.predict(self.frame),True)
        else:
            return block_resultResponse([],False)

    
    def predict(self, frame):
        roi = cv.warpPerspective(frame, self.Mat, (400, 300))
        image = cv.cvtColor(roi, cv.COLOR_BGR2GRAY)
        origin = [0, 0]  # 原点
        length = 80  # 长度
        dst_x = 107  # 偏移值
        dst_y = 72
        result = []
        data = []
        for i in range(4):
            for j in range(4):
                spilt = image[origin[1] + dst_y * i:origin[1] + length + dst_y * i,
                        origin[0] + dst_x * j:origin[0] + length + dst_x * j]
                spilt = cv.resize(spilt, (80, 80))
                data.append(spilt)
        
        data = np.array(data)
        data = data.astype('float32') / 255
        data = data.reshape(data.shape[0], 6400)
        ret, result, neighbors, dist = self.model.findNearest(data, k=3)
        return result.ravel()


if __name__ == "__main__":
    det = Detection()
    det.ros_running()
        